from vocab import *
from typing import Tuple

from torch.utils.data import Dataset, DataLoader

"""
定义了一个机器翻译数据集类MTSet：
init：初始化函数，接收字符串对列表作为数据
getitem：根据索引获取数据项，将源语言和目标语言文本编码为ID序列，并转换为PyTorch张量
len：返回数据集大小
该类用于将文本数据转换为模型可处理的数值张量格式。
"""


class MTSet(Dataset):
    def __init__(self, pairs: List[Tuple[str, str]]):
        self.data = pairs

    def __getitem__(self, index: int):
        src, tgt = self.data[index]
        src_ids = src_vocab.encode(src, add_sos_eos=True)
        tgt_ids = tgt_vocab.encode(tgt, add_sos_eos=True)
        return torch.Tensor(src_ids), torch.Tensor(tgt_ids)

    def __len__(self):
        return len(self.data)


def pad_seq(seq, max_len, pad_id):
    """
    序列填充：
    创建一个长度为max_len的张量，所有元素都填充为pad_id
    将输入序列seq的内容填充到张量的前len(seq)个位置
    返回填充后的张量
    主要用于将不同长度的序列统一填充到相同长度，便于批量处理。
    """
    # Python中元组的创建语法：
    # (max_len,) 表示创建一个包含单个元素 max_len 的元组
    # (max_len) 只是用括号包围的表达式，不是元组
    out = torch.full((max_len,), pad_id, dtype=torch.long)
    out[:len(seq)] = seq
    return out


def collate(batch):
    """
    collate函数用于处理批量数据，将源序列和目标序列分别填充到相同长度，并构造输入输出对。具体步骤如下：
    解包批次数据：将批次中的源序列和目标序列分别提取出来。
    计算最大长度：找出源序列和目标序列中的最大长度，用于后续填充。
    序列填充：使用pad_seq函数将每个序列填充至对应的最大长度，并将结果堆叠成张量。
    构造目标输入输出：将目标序列切分为输入（去掉最后一个词）和输出（去掉第一个词）。
    返回结果：将处理好的张量转移到指定设备并返回。
    """
    src_seqs, tgt_seqs = zip(*batch)
    src_max = max(len(seq) for seq in src_seqs)
    tgt_max = max(len(seq) for seq in tgt_seqs)

    # 对源序列和目标序列分别进行填充，使每个批次内的序列长度一致。
    # 使用 pad_seq 函数将每条序列填充到对应的最大长度（src_max 或 tgt_max），填充值来自词汇表中的 PAD 标记。
    # 使用 torch.stack 将填充后的序列堆叠成一个张量，便于批量处理。
    # 最终得到形状为 (batch_size, max_len) 的张量，用于后续模型输入。
    src_batch = torch.stack([pad_seq(seq, src_max, src_vocab.stoi[PAD]) for seq in src_seqs])
    tgt_batch = torch.stack([pad_seq(seq, tgt_max, tgt_vocab.stoi[PAD]) for seq in tgt_seqs])

    tgt_in = tgt_batch[:, :-1]  # 输入序列，去掉最后一个词
    tgt_out = tgt_batch[:, 1:]  # 输出序列，去掉第一个词
    return src_batch.to(device), tgt_in.to(device), tgt_out.to(device)


"""
DataLoader 是一个可迭代对象，不是列表或数组
每次迭代返回一个批次的数据，由 collate 函数处理生成
返回的数据格式为 (src_batch, tgt_in, tgt_out) 元组
批次大小由 batch_size 参数决定（这里是4）
"""

dataset = MTSet(pairs)
# collate_fn是一个可调用对象（函数），用于处理从数据集中获取的一批数据样本
# 它接收一个列表作为输入，列表中的每个元素是数据集getitem方法返回的内容
# 它需要返回处理后的批次数据
loader = DataLoader(dataset, batch_size=4, shuffle=True, collate_fn=collate)

if __name__ == '__main__':
    # 通过 for 循环遍历 loader
    for batch in loader:
        src_batch, tgt_in, tgt_out = batch
        # 处理批次数据
        print("源序列批次形状:", src_batch.shape)
        print(src_batch)
        print("目标输入批次形状:", tgt_in.shape)
        print(tgt_in)
        print("目标输出批次形状:", tgt_out.shape)
        print(tgt_out)
